// Copyright 2022 Chen Jun

#include "armor_detector/pnp_solver.hpp"


namespace rm_auto_aim
{
PnPSolver::PnPSolver(
  const std::array<double, 9> & camera_matrix, const std::vector<double> & dist_coeffs)
: camera_matrix_(cv::Mat(3, 3, CV_64F, const_cast<double *>(camera_matrix.data())).clone()),
  dist_coeffs_(cv::Mat(1, 5, CV_64F, const_cast<double *>(dist_coeffs.data())).clone())
{
  // Unit: m
  constexpr double small_half_y = SMALL_ARMOR_WIDTH / 2.0 / 1000.0;
  constexpr double small_half_z = SMALL_ARMOR_HEIGHT / 2.0 / 1000.0;
  constexpr double large_half_y = LARGE_ARMOR_WIDTH / 2.0 / 1000.0;
  constexpr double large_half_z = LARGE_ARMOR_HEIGHT / 2.0 / 1000.0;

  // 初始化大小装甲板的装甲板坐标系坐标
  small_armor_points_.emplace_back(cv::Point3f(0, small_half_y, -small_half_z));
  small_armor_points_.emplace_back(cv::Point3f(0, small_half_y, small_half_z));
  small_armor_points_.emplace_back(cv::Point3f(0, -small_half_y, small_half_z));
  small_armor_points_.emplace_back(cv::Point3f(0, -small_half_y, -small_half_z));

  large_armor_points_.emplace_back(cv::Point3f(0, large_half_y, -large_half_z));
  large_armor_points_.emplace_back(cv::Point3f(0, large_half_y, large_half_z));
  large_armor_points_.emplace_back(cv::Point3f(0, -large_half_y, large_half_z));
  large_armor_points_.emplace_back(cv::Point3f(0, -large_half_y, -large_half_z));

  // small_armor_points_.emplace_back(cv::Point3f(-small_half_y, -small_half_z,0));
  // small_armor_points_.emplace_back(cv::Point3f(-small_half_y, small_half_z,0));
  // small_armor_points_.emplace_back(cv::Point3f(small_half_y, small_half_z,0));
  // small_armor_points_.emplace_back(cv::Point3f(small_half_y, -small_half_z,0));

  // large_armor_points_.emplace_back(cv::Point3f(-large_half_y, -large_half_z,0));
  // large_armor_points_.emplace_back(cv::Point3f(-large_half_y, large_half_z,0));
  // large_armor_points_.emplace_back(cv::Point3f(large_half_y, large_half_z,0));
  // large_armor_points_.emplace_back(cv::Point3f(large_half_y, -large_half_z,0));
}

bool PnPSolver::solvePnP(const Armor & armor, cv::Mat & rvec, cv::Mat & tvec)
{
  //装甲板在相机图像中的坐标，图像左上，左下，右上，右下
  std::vector<cv::Point2f> image_armor_points;

  image_armor_points.emplace_back(armor.armor_verticles_[0]);
  image_armor_points.emplace_back(armor.armor_verticles_[1]);
  image_armor_points.emplace_back(armor.armor_verticles_[2]);
  image_armor_points.emplace_back(armor.armor_verticles_[3]);

  // PNP解算
  auto object_points = armor.type == ArmorType::SMALL ? small_armor_points_ : large_armor_points_;
  bool success = cv::solvePnP(
    object_points, image_armor_points, camera_matrix_, dist_coeffs_, rvec, tvec, false,
    cv::SOLVEPNP_IPPE);
  if (!success) {
    return false;
  }

  // // 获取yaw
  // double yaw = goldenSreach(-80, 80, 1, image_armor_points, tvec, object_points);
  // cv::Mat rotation_matrix;
  // // 旋转矩阵转换为旋转向量以修正yaw
  // cv::Rodrigues(rvec, rotation_matrix);
  // //修正yaw
  // rotation_matrix.at<double>(0, 2) = yaw;
  // cv::Rodrigues(rotation_matrix, rvec);
  return true;
}

float PnPSolver::calculateDistanceToCenter(const cv::Point2f & image_point)
{
  float cx = camera_matrix_.at<double>(0, 2);
  float cy = camera_matrix_.at<double>(1, 2);
  return cv::norm(image_point - cv::Point2f(cx, cy));
}

cv::Point2f PnPSolver::reproject(Eigen::Vector3d & xyz)
{
  Eigen::Matrix3d mat_intrinsic;
  cv::cv2eigen(camera_matrix_, mat_intrinsic);
  //(u,v,1)^T = (1/Z) * K * (X,Y,Z)^T
  auto result = (1.f / xyz[2]) * mat_intrinsic * (xyz);  //解算前进行单位转换
  return cv::Point2f(result[0], result[1]);
}

cv::Mat euler2Rotation( double yaw, double pitch, double roll )
{
    // Convert degrees to radians
    yaw = yaw * CV_PI / 180.0;
    pitch = pitch * CV_PI / 180.0;
    roll = roll * CV_PI / 180.0;

    cv::Mat rotation = cv::Mat::eye( 3, 3, CV_64F );

    // Z-Y-X rotation order
    double sy = sin( yaw );
    double cy = cos( yaw );
    double sp = sin( pitch );
    double cp = cos( pitch );
    double sr = sin( roll );
    double cr = cos( roll );

    rotation.at< double >( 0, 0 ) = cy * cp;
    rotation.at< double >( 0, 1 ) = cy * sp * sr - sy * cr;
    rotation.at< double >( 0, 2 ) = cy * sp * cr + sy * sr;
    rotation.at< double >( 1, 0 ) = sy * cp;
    rotation.at< double >( 1, 1 ) = sy * sp * sr + cy * cr;
    rotation.at< double >( 1, 2 ) = sy * sp * cr - cy * sr;
    rotation.at< double >( 2, 0 ) = -sp;
    rotation.at< double >( 2, 1 ) = cp * sr;
    rotation.at< double >( 2, 2 ) = cp * cr;

    return rotation;
}


//tip: reproject重投影到的是矫正畸变后的图像，而cost是和原图上的比较，是否会有误差
//tip: pitch固定，在机器人抬头低头过程中，因为pitch不真实，是否会存在误差

double PnPSolver::countRemapYawCost(double yaw, cv::Mat tvec,
  std::vector<cv::Point3f> object_points, std::vector<cv::Point2f> image_armor_points)
{
  // 矩形在坐标系下的yaw、pitch、roll
  double pitch = 15.0;
  double roll = 0.0;

  // 相机坐标系下装甲板
  double X_c = tvec.at<double>(0);
  double Y_c = tvec.at<double>(1);
  double Z_c = tvec.at<double>(2);

  cv::Mat rotation_matrix;
  Rodrigues(cv::Vec3f(CV_PI / 180 * pitch, CV_PI / 180 * yaw, CV_PI / 180 * roll), rotation_matrix);

  // 旋转后的坐标
  std::vector<cv::Point3f> points_rotated;

  // 应用旋转矩阵
  transform(
    object_points.begin(), object_points.end(), std::back_inserter(points_rotated),
    [&](const cv::Point3f & point) {
      cv::Mat point_local_mat = (cv::Mat_<float>(3, 1) << point.x, point.y, point.z);
      cv::Mat point_rotated_mat = rotation_matrix * point_local_mat;
      cv::Point3f point_rotated(
        point_rotated_mat.at<float>(0), point_rotated_mat.at<float>(1),
        point_rotated_mat.at<float>(2));
      return point_rotated;
    });

  std::vector<Eigen::Vector3d> points_world = {
    Eigen::Vector3d(X_c + points_rotated[0].x, Y_c + points_rotated[0].y, Z_c + points_rotated[0].z),
    Eigen::Vector3d(X_c + points_rotated[1].x, Y_c + points_rotated[1].y, Z_c + points_rotated[1].z),
    Eigen::Vector3d(X_c + points_rotated[2].x, Y_c + points_rotated[2].y, Z_c + points_rotated[2].z),
    Eigen::Vector3d(X_c + points_rotated[3].x, Y_c + points_rotated[3].y, Z_c + points_rotated[3].z)};
  std::vector<cv::Point2f> points_camera = {
    reproject(points_world[0]), reproject(points_world[1]), reproject(points_world[2]),
    reproject(points_world[3])};

  // 计算目标函数值，使用二次型形式
  std::vector<cv::Point2f> difference;
  std::transform(
    points_camera.begin(), points_camera.end(), image_armor_points.begin(),
    std::back_inserter(difference), [](const cv::Point2d & p1, const cv::Point2d & p2) { return p1 - p2; });
  double objective_value = 0.0;
  for (const auto & point_diff : difference) {
    objective_value += point_diff.ddot(point_diff);
  }

  return objective_value;
}

double PnPSolver::goldenSreach(
  double left, double right, double epsilon, std::vector<cv::Point2f> & image_armor_points,
  cv::Mat tvec, std::vector<cv::Point3f> object_points)
{
  const double phi = (1.0 + sqrt(5.0)) / 2.0;
  double x1 = right - (right - left) / phi;
  double x2 = left + (right - left) / phi;

  double f1 = countRemapYawCost(x1, tvec, object_points,image_armor_points);
  double f2 = countRemapYawCost(x2, tvec, object_points,image_armor_points);

  while ((right - left) > epsilon) {
    // i++;
    if (f1 < f2) {
      right = x2;
      x2 = x1;
      x1 = right - (right - left) / phi;
      f2 = f1;
      f1 = countRemapYawCost(x1, tvec, object_points,image_armor_points);
    } else {
      left = x1;
      x1 = x2;
      x2 = left + (right - left) / phi;
      f1 = f2;
      f2 = countRemapYawCost(x2, tvec, object_points,image_armor_points);
    }
  }
  return (left + right) / 2;
}

}  // namespace rm_auto_aim
